#!/usr/bin/env python3
"""
LangChain MCP 智能体调用框架
功能：
1. 连接多个MCP服务器
2. 自动发现和注册MCP工具
3. 创建带记忆的智能体
4. 支持自然语言交互和工具选择
"""

import os
import asyncio
from typing import Dict, Any, List
from dotenv import load_dotenv

# LangChain 核心模块
from langchain_core.tools import Tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory

# 语言模型
from langchain_community.llms import Tongyi

# 智能体相关
from langchain.agents import create_react_agent, AgentExecutor

# MCP 相关
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

# 加载环境变量
load_dotenv()

# ==================== MCP 智能体类 ====================

class MCPAgent:
    """基于MCP的LangChain智能体框架类"""
    
    def __init__(self):
        self.llm = None
        self.tools = []
        self.agent_executor = None
        self.store = {}  # 会话存储
        self.mcp_sessions = {}  # MCP会话存储
        
    def setup_llm(self):
        """初始化语言模型"""
        print("正在初始化语言模型...")
        
        api_key = os.getenv("DASHSCOPE_API_KEY")
        if not api_key:
            raise ValueError("请在.env文件中设置DASHSCOPE_API_KEY")
            
        self.llm = Tongyi(
            model="qwen-turbo",
            api_key=api_key,
            temperature=0.7,
            max_tokens=1024
        )
        print("语言模型初始化完成")
        
    async def setup_mcp_servers(self):
        """设置MCP服务器连接"""
        print("正在连接MCP服务器...")
        
        # 定义MCP服务器配置
        mcp_servers = {
            "filesystem": {
                "command": "npx",
                "args": ["-y", "@modelcontextprotocol/server-filesystem", "."],
                "env": {}
            },
            "brave-search": {
                "command": "npx",
                "args": ["-y", "@modelcontextprotocol/server-brave-search"],
                "env": {
                    "BRAVE_API_KEY": os.getenv("BRAVE_API_KEY", "")
                }
            },
            "memory": {
                "command": "npx",
                "args": ["-y", "@modelcontextprotocol/server-memory"],
                "env": {}
            }
        }
        
        # 连接每个MCP服务器
        for server_name, config in mcp_servers.items():
            try:
                print(f"连接到 {server_name} 服务器...")
                
                # 创建服务器参数
                server_params = StdioServerParameters(
                    command=config["command"],
                    args=config["args"],
                    env=config.get("env", {})
                )
                
                # 建立连接
                stdio_transport = await stdio_client(server_params)
                stdio, write = stdio_transport
                
                # 创建会话
                session = ClientSession(stdio, write)
                await session.initialize()
                
                # 存储会话
                self.mcp_sessions[server_name] = session
                
                # 获取工具列表
                tools_result = await session.list_tools()
                
                # 为每个工具创建LangChain Tool
                for tool_info in tools_result.tools:
                    langchain_tool = Tool(
                        name=f"{server_name}_{tool_info.name}",
                        description=tool_info.description or f"{server_name}工具: {tool_info.name}",
                        func=self.create_mcp_tool_wrapper(server_name, tool_info.name)
                    )
                    self.tools.append(langchain_tool)
                    
                print(f"{server_name} 服务器连接成功，发现 {len(tools_result.tools)} 个工具")
                
            except Exception as e:
                print(f"连接 {server_name} 服务器失败: {e}")
                # 创建模拟工具以确保框架正常运行
                mock_tool = Tool(
                    name=f"{server_name}_mock",
                    description=f"模拟的{server_name}工具",
                    func=lambda x: f"模拟{server_name}工具执行: {x}"
                )
                self.tools.append(mock_tool)
                
        print(f"MCP服务器设置完成，共发现 {len(self.tools)} 个工具")
        
    def create_mcp_tool_wrapper(self, server_name: str, tool_name: str):
        """创建MCP工具包装器"""
        async def mcp_tool_func(input_data: str) -> str:
            try:
                session = self.mcp_sessions.get(server_name)
                if not session:
                    return f"错误: {server_name} 服务器未连接"
                
                # 解析输入参数
                try:
                    import json
                    if isinstance(input_data, str):
                        if input_data.startswith('{'):
                            arguments = json.loads(input_data)
                        else:
                            arguments = {"input": input_data}
                    else:
                        arguments = input_data
                except:
                    arguments = {"input": str(input_data)}
                
                # 调用MCP工具
                result = await session.call_tool(tool_name, arguments)
                
                # 处理结果
                if result.content:
                    return str(result.content[0].text if result.content else "无结果")
                else:
                    return "工具执行完成，无返回内容"
                    
            except Exception as e:
                return f"调用{server_name}.{tool_name}时出错: {str(e)}"
        
        # 创建同步包装器
        def sync_wrapper(input_data: str) -> str:
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    # 如果已经在事件循环中，创建新的任务
                    import concurrent.futures
                    with concurrent.futures.ThreadPoolExecutor() as executor:
                        future = executor.submit(asyncio.run, mcp_tool_func(input_data))
                        return future.result()
                else:
                    return asyncio.run(mcp_tool_func(input_data))
            except Exception as e:
                return f"工具执行错误: {str(e)}"
        
        return sync_wrapper
        
    def setup_agent(self):
        """创建智能体"""
        print("正在创建智能体...")
        
        # 创建标准的ReAct提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个智能助手，可以使用多种MCP工具来帮助用户。"),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", """
你有以下工具可以使用:
{tools}

使用以下格式:

Question: 用户的问题
Thought: 你应该思考要做什么
Action: 要采取的行动，应该是 [{tool_names}] 中的一个
Action Input: 行动的输入
Observation: 行动的结果
... (这个 Thought/Action/Action Input/Observation 可以重复N次)
Thought: 我现在知道最终答案了
Final Answer: 对原始问题的最终答案

开始!

Question: {input}
Thought:{agent_scratchpad}""")
        ])
        
        # 创建基础智能体
        agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=prompt
        )
        
        # 创建智能体执行器
        base_executor = AgentExecutor(
            agent=agent,
            tools=self.tools,
            verbose=True,
            handle_parsing_errors=True,
            max_iterations=5
        )
        
        # 添加记忆功能
        self.agent_executor = RunnableWithMessageHistory(
            base_executor,
            self.get_session_history,
            input_messages_key="input",
            history_messages_key="chat_history",
        )
        
        print("智能体创建完成")
        
    def get_session_history(self, session_id: str) -> InMemoryChatMessageHistory:
        """获取会话历史"""
        if session_id not in self.store:
            self.store[session_id] = InMemoryChatMessageHistory()
        return self.store[session_id]
        
    async def initialize(self):
        """初始化智能体"""
        print("=== 初始化 MCP 智能体框架 ===")
        self.setup_llm()
        await self.setup_mcp_servers()
        self.setup_agent()
        print("\n=== 初始化完成 ===")
        print(f"可用工具: {[tool.name for tool in self.tools]}")
        
    def chat(self, user_input: str, session_id: str = "default") -> str:
        """与智能体对话"""
        try:
            response = self.agent_executor.invoke(
                {"input": user_input},
                config={"configurable": {"session_id": session_id}}
            )
            return response["output"]
        except Exception as e:
            return f"处理请求时出错: {str(e)}"
            
    async def interactive_chat(self):
        """交互式对话"""
        print("\n=== 开始交互式对话 ===")
        print("输入 'quit' 或 'exit' 退出对话")
        print("输入 'clear' 清除对话历史")
        print("-" * 50)
        
        session_id = "user_session"
        
        while True:
            try:
                user_input = input("\n用户: ").strip()
                
                if user_input.lower() in ['quit', 'exit', '退出']:
                    print("再见！")
                    break
                    
                if user_input.lower() in ['clear', '清除']:
                    if session_id in self.store:
                        del self.store[session_id]
                    print("对话历史已清除")
                    continue
                    
                if not user_input:
                    continue
                    
                print("\n助手: ", end="")
                response = self.chat(user_input, session_id)
                print(response)
                
            except KeyboardInterrupt:
                print("\n\n对话已中断")
                break
            except Exception as e:
                print(f"\n发生错误: {e}")
                
    async def cleanup(self):
        """清理资源"""
        print("正在清理MCP连接...")
        for server_name, session in self.mcp_sessions.items():
            try:
                await session.close()
                print(f"{server_name} 连接已关闭")
            except Exception as e:
                print(f"关闭 {server_name} 连接时出错: {e}")

# ==================== 测试函数 ====================

async def test_mcp_agent():
    """测试MCP智能体功能"""
    print("=== 开始测试 ===")
    
    # 创建智能体实例
    agent = MCPAgent()
    await agent.initialize()
    
    # 测试对话
    test_cases = [
        "你好，我是小明",
        "帮我搜索一下人工智能的最新发展",
        "请列出当前目录下的文件",
        "帮我记住：我喜欢喝咖啡",
        "你还记得我刚才说的喜好吗？"
    ]
    
    session_id = "test_session"
    
    for i, test_input in enumerate(test_cases, 1):
        print(f"\n--- 测试 {i} ---")
        print(f"用户: {test_input}")
        response = agent.chat(test_input, session_id)
        print(f"助手: {response}")
        print("-" * 30)
    
    return agent

# ==================== 主函数 ====================

async def main():
    """主函数"""
    agent = None
    try:
        # 运行测试
        agent = await test_mcp_agent()
        
        # 启动交互式对话
        await agent.interactive_chat()
        
    except Exception as e:
        print(f"程序运行出错: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # 清理资源
        if agent:
            await agent.cleanup()

if __name__ == "__main__":
    asyncio.run(main())